chore(sdk): upgrade to Neuron SDK 2.31 - #1113

The SDK version is part of the test model hub repository names, so it must
be updated along with the dependencies to avoid reusing artifacts compiled
with the previous compiler.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

torch_neuronx keeps a process-global HLO instance counter that gets
baked into the compiled module hash, so exporting the same model
twice in one process always produces different hashes and the
Hub cache lookup never hits. Run each export in its own subprocess.

_sample computed top_k=1 for greedy, but built the CPU logits warper
from the raw generation_config, so greedy decoding actually ran
topk(logits, 50) and took element 0. On an exact bfloat16 tie that
breaks towards whatever the topk sort returns instead of the lowest
token id, which is what transformers does. token_selector.select() and
_assisted_decoding already used a plain argmax, so the same model could
disagree with itself between speculative and regular greedy decoding.

Select from the full logits when do_sample is False, and keep the fused
warper on the sampling path only. This is host side: the graph inputs
are unchanged and sampling_params stays (batch_size, 3) whatever the
sampling parameters are, so no recompilation is triggered.

qwen2-4x1024 and qwen2-1x8192 now reproduce the CPU output exactly, and
granite-4x1024 and qwen3-1x8192 turned out to already match, so their
entries were stale. The greedy expectations go from 9 passed/4 xfailed
to 11 passed/2 xfailed.

qwen3-4x1024 and qwen3-tp1-4x1024 still differ, and no modeling change
can help: CPU resolves that token with a 0.0317 logit gap, a quarter of
a bfloat16 ULP at that magnitude (0.125). Both now generate what the
CPU model generates in bfloat16.

The gemma3 long-sequence xfail is kept as is, but its margin is 3.44
ULP, so unlike the others it is not a rounding tie and still needs an
explanation.

Also run test_speculation_same_model in a subprocess: loading a
compiled model onto NeuronCores sets the Neuron runtime world_size for
the process lifetime, so the preceding tests in the file make it fail
with "Could not load the model" when the whole file is run.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Neuron picks CPU's second-best candidate at the third generated
token (17.10 vs 16.93 logits, a near-tie), producing coherent but
different text for smolvlm-16-images-cross-chunk. Xfail on the
known observed string, same pattern as the decoder greedy tests.

Moves import paths (FlexibleArgumentParser, WorkerBase,
set_random_seed), ModelConfig.task -> runner_type, and the CLI's
--task -> --runner/--convert split. Also fixes behavioral breaks
the import-only survey missed: WorkerBase.__init__ now requires
local_rank/rank/distributed_init_method directly, load_model must
run inside set_current_vllm_config (CustomOp instantiation asserts
on it), SchedulerConfig lost max_model_len (moved to ModelConfig),
and 0.16 defaults to async scheduling, which needs an execute_model/
sample_tokens split our worker doesn't implement -- disabled it.

pyproject.toml's vllm pin stays uncommitted for now to avoid paying
the full-suite re-export cost more than once.

Bump CI's install_neuronx_runtime action and the vLLM Dockerfile to
the SDK 2.31 DLC runtime versions (tools/runtime-lib/collectives,
plus dkms in the Dockerfile only -- the CI runner gets its driver
from the AMI). Verified via a real image build:

- Ubuntu 22.04 ships Python 3.10; optimum-neuron needs >=3.11 since
  SDK 2.31 ships no cp310 wheels. Provision Python 3.12 via uv (not
  3.11: neuronx-cc pins numpy<2 there, clashing with vllm's numpy>=2).
- torch_xla's compiled extension needs libpython on the linker path,
  which uv's standalone interpreter ships but doesn't register.
- neuronx-cc's compiler binary now needs libarchive13, not in the
  base image.

Also add .dockerignore: with no venv/.git exclusions, docker build
was sending the whole repo (7GB+ locally) as build context, which
likely also slows down CI's vllm-docker-tests job since setup_venv
creates a venv in the repo root before the image build runs.

Also ignore tests/PostSPMDPassesExecutionDuration.txt, a new compiler
artifact under SDK 2.31.

vllm==0.11.0 requires torch==2.8.0, unresolvable against the SDK
2.31 torch==2.9.1 pin, so every vLLM CI job was failing at install.
0.16.0 (and the 0.14-0.16 range) pins torch==2.9.1/torchvision==0.24.1,
matching what's already required -- no torch change needed. This is
the only pyproject.toml edit in this migration, landed last since it
invalidates the decoder test-model re-export cache (~1h, once).

The large-d NKI flash attention kernel stored the softmax numerator in
bfloat16 before the PV matmul, quantizing every attention weight to 8
mantissa bits. That error accumulates over attended positions, which is
why it only ever hit a long prompt: under SDK 2.31 it was enough to
change a sampled token at ~5k tokens, and the long sequence test was
marked as a known divergence from the CPU float32 reference.

Keep p_local, its reduction output and p_local_transposed in fp32 so
full precision reaches the PV matmul. Measured on gemma3-270m at
batch 1 / sequence 8192, first divergence over 50 greedy tokens and
median prefill latency:

  bfloat16 numerator     step 2 (3.44 ULP)       142 ms
  fp32 numerator only    step 2 (bit identical)  161 ms
  flash kernel disabled  step 20 (1.15 ULP)      317 ms
  fp32 through matmul    no divergence           222 ms

Keeping only the matmul operand in bfloat16 reproduces the original
output bit for bit, which localizes the loss to that operand rather
than to the reduction or the softmax denominator. Prefill is ~1.6x
slower than before and still ~1.4x faster than the compiler-native
path. SBUF grows by ~512 KB, independent of head_dim.

Only models reaching the large-d kernel are affected: head_dim > 128
with a prefill of at least 4096 (LNC1) or 2048 (LNC2).

The generation now matches the reference, so the long sequence test no
longer needs its known-divergence escape hatch.

check_and_update_config assigned a closure to
model_config.verify_with_parallel_config. Closures cannot be pickled, so
when vLLM spawns the EngineCore process it fails with "Can't pickle local
object". Replace it with a module-level function so the patched config
survives pickling.

vLLM's V1 engine forks an EngineCore process by default
(VLLM_WORKER_MULTIPROC_METHOD=fork). Forking after torch and the Neuron
runtime have initialized their native thread pools leaves the child with a
dead neuron::ThreadPool, so weight loading deadlocks in
neuron::parallel_load (pthread_barrier_wait) or aborts with "Invalid
thread pool!". Force spawn in the platform plugin's register() so every
vLLM-on-Neuron usage starts EngineCore from a clean interpreter.
setdefault respects an explicit user override. This mirrors what
optimum-cli neuron serve already does.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

convbert, hubert, wav2vec2 and yolos segfault/abort inside
torch_neuronx HLO generation with Neuron SDK 2.31 (torch-neuronx 2.9 /
torch-xla 2.9). The crash is in the compiler's tracer, not in
optimum-neuron code, so it cannot be caught and kills the whole pytest
process. Skip these model types before export until the SDK is fixed.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

The SDK 2.31 DLC is based on Ubuntu 24.04, and the 22.04 runners no
longer match the environment the packages are built for.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

The SDK 2.31 neuronx-cc only ships cp311/cp312/cp313 wheels, and
Ubuntu 24.04 runners ship system Python 3.12. Pin the setup_venv and
sanity-check venvs to 3.12.

Match the CPU torch wheel with the SDK's torch 2.9.1 (previously 2.8).
The CPU wheel sanity-check is cached against these package versions.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

neuronx-cc's weight-layout-optimization step compiles NKI kernels via
torch_neuronx.xla_impl.trace.hlo_compile, which runs the compiler with
subprocess.run(command) and no cwd= (trace.py). The walrus backend then
materializes content-addressed neuronxcc.private_nkl.* kernel dirs in the
process CWD, dirtying the caller's working directory on every cache-miss
export.

Wrap builder.trace() in a context manager that chdirs to a throwaway
temp dir during compilation so those droppings are discarded.

Verified on inf2: exporting katuni4ka/tiny-random-phi3 (fp32, cold NKI
cache) without the fix leaks a 16-hex kernel dir into the CWD; with the
fix the same fresh compile leaves the CWD clean.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

The Ubuntu 24.04 inf2 runner lacks two shared libs SDK 2.31 needs:
libpython3.12.so.1.0 (torch_xla._XLAC import) and libarchive.so.13
(neuronxcc walrus_driver). Verified green on run 32349690521.

Co-Authored-By: Claude <noreply@anthropic.com>

The list of model types crashing the Neuron SDK 2.31 tracer lived in
`tests/exporters/test_transformers.py`, but the same crash affects other
test suites. Move it next to the other test helpers and expose a
`skip_if_sdk_231_trace_crash()` guard so every suite can reuse it
instead of duplicating the list.

No behaviour change: the exporters tests skip exactly the same models.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The conv-based models that crash the Neuron SDK 2.31 tracer were only
skipped in the exporters tests, so the transformers inference and
pipeline suites still exported them and died with SIGSEGV, taking the
whole pytest process down:

- `pytest tests/pipelines` -> convbert, in `F.unfold`
- `pytest -m "not slow" tests/inference/transformers/test_modeling.py`
  -> yolos, in `F.interpolate`
- `pytest -m slow tests/inference/transformers/test_modeling.py`
  -> convbert, in `F.unfold`

Guard the shared `NeuronModelTestMixin._setup()`, which every modeling
test goes through, and the `inf_encoder_model` fixture used by the
pipeline tests. The latter is now parametrized on the architecture
instead of the model id so the skip can be keyed on it.

Both suites pass on main, i.e. with SDK 2.30, so only the tracer version
changed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The `vllm_docker_launcher` context manager stopped and removed the
container after its `yield`, without a `try`/`finally`. When a test
raised inside the `with` block, the exception was thrown back in at the
yield and the whole teardown was skipped, leaving the container running
and holding its Neuron cores.

The next test using the same device then died during startup:

    NRT:nrt_allocate_neuron_cores  Logical Neuron Core(s) not available
    - Requested:lnc8-lnc9 Available:0 Logical Core size:1
    (cores busy, ret=-16)
    RuntimeError: The PyTorch Neuron Runtime could not be initialized.

So a single test failure cascaded into the next one. Wrap the yield so
the container is always stopped and removed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

`test_vllm_docker_service_sampling_parameters` asserts that a sampled
answer differs from the greedy one. The output distribution of
Llama-3.2-1B on this prompt is peaked enough that a single draw at
temperature 1.0 / top_p 0.9 sometimes reproduces the greedy answer, so
the assertion is flaky. Bumping the temperature from 0.8 to 1.0 in
7d9d1ac made it rarer but did not remove it.

Draw up to five samples and stop as soon as one differs, which is what
the test actually means to check: that the sampling parameters are taken
into account. A single draw matching greedy is not a failure.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>